﻿namespace SpikingAnalyticsLib

open System.IO
open SpikingNeuronLib
open IOUtilities

// ===============================================================================================================
// Utility functions for extracting graphs from peak response data
// Define 'MapPrePost' as a Map of the form: preNeuron * ((postNeuron, peakTime, peakCount) list)

module GraphUtilities =

    // Read the data in the specified file into a map
    let LoadMapPrePost filePath =
        let _, data =
            File.ReadLines(filePath)
            |> Seq.fold (fun (state:int * Map<int, (int * int * int) list>) line ->
                    let key, acc = state
                    if line.Contains(" ") then
                        let tokenList = Tokenize (List.ofArray (line.ToCharArray()))
                        let data = ConstructTuple<int * int * int>(tokenList)
                        let responseList =
                            if acc.ContainsKey(key) then
                                acc.[key]
                            else
                                []
                        key, Map.add key (data :: responseList) acc
                    else
                        let nextKey = System.Int32.Parse(line)
                        nextKey, acc
                ) (0, Map.empty)
        data

    // Get the list of neurons in the longest path starting from the specified neuron using data in the specified MapPrePost
    let GetLongestPath preNeuron mapPrePost =
        let rec GetLongestPath preNeuron (visited:Set<int>) (mapPrePost:Map<int, (int * int * int) list>) : int list =
            let longestPath =
                if (not (visited.Contains(preNeuron))) && (mapPrePost.ContainsKey preNeuron) then
                    mapPrePost.[preNeuron]
                    |> Seq.map (fun (postNeuron, _, _) ->
                            GetLongestPath postNeuron (visited.Add preNeuron) mapPrePost
                        )
                    |> Seq.maxBy (fun path -> path.Length)
                else
                    []
            preNeuron :: longestPath
        GetLongestPath preNeuron (Set.empty) mapPrePost

    // Get the path length (number of hops) starting from the specified neuron using data in the specified MapPrePost
    let GetLongestPathLength preNeuron mapPrePost =
        let rec GetLongestPath preNeuron pathLength (visited:Set<int>) (mapPrePost:Map<int, (int * int * int) list>) : int =
            let maxLength =
                if (not (visited.Contains(preNeuron))) && (mapPrePost.ContainsKey preNeuron) then
                    mapPrePost.[preNeuron]
                    |> Seq.map (fun (postNeuron, _, _) ->
                            GetLongestPath postNeuron (pathLength + 1) (visited.Add preNeuron) mapPrePost
                        )
                    |> Seq.max
                else
                    pathLength
            maxLength
        GetLongestPath preNeuron 0 (Set.empty) mapPrePost

    // Create a mapping between neurons and their firing (event) times using data in the specified MapPrePost
    let CreateEventTimeMap downstreamMap =

        // Given a map from pre to multiple post neurons, create a map from post to multiple pre neurons
        let ReverseMapData mapPrePost =
            mapPrePost
            |> Map.toSeq
            |> Seq.map (fun (preNeuron, postList) ->
                    postList
                    |> Seq.map (fun (postNeuron, peakTime, peakCount) -> (preNeuron, (postNeuron, peakTime, peakCount)))
                )
            |> Seq.concat
            |> Seq.fold (fun (state:Map<int, Map<int, bool>>) (preNeuron, (postNeuron, _, _)) ->
                    let preMap =
                        if state.ContainsKey postNeuron then
                            state.[postNeuron]
                        else
                            Map.empty
                    Map.add postNeuron (Map.add preNeuron false preMap) state
                ) Map.empty

        // Produce a sequence of upstream neurons
        let GetUpStreamBranch postNeuron upstreamMap =
            let rec GetUpStreamBranch neuron (visited:Set<int>) (upstreamMap:Map<int, Map<int, bool>>) =
                seq {
                    if (not (visited.Contains(neuron))) && (upstreamMap.ContainsKey neuron) then
                        let preMapSequence = Map.toSeq upstreamMap.[neuron]
                        for (preNeuron, connected) in preMapSequence do
                            if connected then
                                yield preNeuron
                                yield! GetUpStreamBranch preNeuron (visited.Add neuron) upstreamMap
                    }
            GetUpStreamBranch postNeuron Set.empty upstreamMap |> Set.ofSeq

        // Get the set of all neurons with a known connection to the specified neuron
        let GetConnectedNeurons postNeuron (downstreamMap:Map<int, (int * int * int) list>) (upstreamMap:Map<int, Map<int, bool>>) =
            let upstreamNeurons = GetUpStreamBranch postNeuron upstreamMap
            let downstreamNeurons =
                if downstreamMap.ContainsKey postNeuron then
                    downstreamMap.[postNeuron]
                    |> Seq.map (fun (postNeuron, _, _) -> postNeuron)
                    |> Set.ofSeq
                else
                    Set.empty
            let additionalNeurons = // upstream from the downstream neurons
                if downstreamNeurons.IsEmpty then
                    Set.empty
                else
                    downstreamNeurons
                    |> Seq.map (fun neuron ->
                            GetUpStreamBranch neuron upstreamMap
                        )
                    |> Set.unionMany

            Set.unionMany (seq
                {
                    yield upstreamNeurons
                    yield downstreamNeurons
                    yield additionalNeurons
                    yield Set.singleton postNeuron
                })

        // Produce a new map where the connection flag is toggled on for the specified pre- and post-synaptic neurons
        let ConnectNeurons preNeuron postNeuron (upstreamMap:Map<'a, Map<'b, bool>>) =
            let preMap = Map.add preNeuron true upstreamMap.[postNeuron]
            Map.add postNeuron preMap upstreamMap

        // Adjust all times on a mapping between neuron and firing time so that the earliest time maps to time t=0
        let NormaliseEventMap (map:Map<int, int>) =
            let _, minTime =
                map
                |> Map.toSeq
                |> Seq.minBy (fun (neuron, time) -> time)
            map
            |> Map.map (fun neuron time -> time - minTime)

        let upstreamMap = ReverseMapData downstreamMap
        let firingEventMap, _ =
            downstreamMap
            |> Map.toSeq
            |> Seq.map (fun (preNeuron, postList) ->
                    postList
                    |> Seq.map (fun (postNeuron, peakTime, peakCount) -> preNeuron, (postNeuron, peakTime, peakCount))
                )
            |> Seq.concat
            |> Seq.fold (fun (state:Map<int, int> * Map<int, Map<int, bool>>) (preNeuron, (postNeuron, peakTime, _)) ->
                    let firingEventMap, upstreamMap = state
                    let updatedFiringEventMap =
                        if firingEventMap.ContainsKey(postNeuron) then
                            if firingEventMap.ContainsKey(preNeuron) then
                                let adjustment = (firingEventMap.[preNeuron] + peakTime) - firingEventMap.[postNeuron]
                                // adjust the event timing of the entire branch of connected neurons
                                let connectedNeurons = GetConnectedNeurons postNeuron downstreamMap upstreamMap
                                connectedNeurons
                                |> Seq.fold (fun (map:Map<int, int>) connectedNeuron ->
                                        if map.ContainsKey(connectedNeuron) then
                                            let adjustedTime = map.[connectedNeuron] + adjustment
                                            Map.add connectedNeuron adjustedTime map
                                        else
                                            map
                                    ) firingEventMap
                            else
                                let time = firingEventMap.[postNeuron]
                                Map.add preNeuron (time - peakTime) firingEventMap
                        else
                            let temp = Map.add preNeuron 0 firingEventMap
                            Map.add postNeuron peakTime temp
                    let updatedUpstreamMap = ConnectNeurons preNeuron postNeuron upstreamMap
                    updatedFiringEventMap, updatedUpstreamMap
                ) (Map.empty, upstreamMap)

        NormaliseEventMap firingEventMap

    // Get any subgraphs in the specified MapPrePost
    let GetGroups maxTime minCount data =

        // Get a list view of the tree of firing events that are downstream from the specified neuron
        let GetLinkedEvents preNeuron eventTimeMap mapPrePost =
            let rec GetLinkedEvents preNeuron (visited:Set<int>) (eventTimeMap:Map<int, int>) (mapPrePost:Map<int, (int * int * int) list>) =
                let eventList =
                    seq {
                        if (not (visited.Contains(preNeuron))) && (mapPrePost.ContainsKey preNeuron) then
                            for (postNeuron, peakTime, peakCount) in mapPrePost.[preNeuron] do
                                // if the time difference < 20 msec
                                if eventTimeMap.[postNeuron] - eventTimeMap.[preNeuron] < 20 then
                                    // connection weight of 0.0 and link probability of 1.0 added for compatibility w/ new LinkedEventPair version
                                    yield (eventTimeMap.[preNeuron], preNeuron), (eventTimeMap.[postNeuron], postNeuron), 0.0, 1.0
                                    yield! GetLinkedEvents postNeuron (visited.Add preNeuron) eventTimeMap mapPrePost
                        }
                eventList
            GetLinkedEvents preNeuron (Set.empty) eventTimeMap mapPrePost |> Set.ofSeq

        // Produce a new map that is a filtered copy of the specified MapPrePost
        let FilterMap (maxTime:int option) (minCount:int option) mapPrePost =
            mapPrePost
            |> Map.toSeq
            |> Seq.map (fun (preNeuron, postList) ->
                    let filteredList =
                        postList
                        |> Seq.filter (fun (postNeuron, peakTime, peakCount) -> maxTime.IsNone || (maxTime.IsSome && peakTime <= maxTime.Value))
                        |> Seq.filter (fun (postNeuron, peakTime, peakCount) -> minCount.IsNone || (minCount.IsSome && peakCount >= minCount.Value))
                        |> Seq.toList
                    preNeuron, filteredList
                )
            |> Map.ofSeq

        let filteredData = FilterMap maxTime minCount data
        let eventTimeMap = CreateEventTimeMap filteredData  // use firing times specified in this firing event map
        let linkedEventsList =
            filteredData
            |> Map.toSeq
            |> Seq.map (fun (preNeuron, _) -> GetLinkedEvents preNeuron eventTimeMap filteredData)
            |> Set.unionMany
            |> Set.toList

        let minTime, maxTime =
            let timeList =
                eventTimeMap
                |> Map.toSeq
                |> Seq.map (fun (neuron, firingTime) -> firingTime)
                |> Seq.toList
            Seq.min timeList, Seq.max timeList

        minTime, maxTime, linkedEventsList

    // Filter (or search for) firing events using the event pairs in the specified MapPrePost
    // Selects events that are part of an establised relationship and
    // therefore probably part of the same group
    // Note: will search for multiple groups simultaneously
    let FilterEvents jitter (mapPrePost:Map<int, _>) firingEvents =

        let CleanMap event jitter watchMap =
            let time, _, _ = event
            let maxTime = time - jitter - 10
            watchMap
            |> Map.filter (fun key _ -> key >= maxTime)

        let SelectWatchedEvents (mapPrePost:Map<int, _>) preEvent currentWatchMap =
            let time, nindex, _ = preEvent
            if mapPrePost.ContainsKey nindex then
                // for each pre-post pair where pre = nindex, add post to the watchMap with appropriate delay
                mapPrePost.[nindex]
                |> Seq.fold (fun (map:Map<int, Map<int * int, (int * int * EventLabel) list>>) (postNeuron, peakTime, peakCount) ->
                        let watchTime = (time + peakTime)
                        let postEvent = (watchTime, postNeuron)
                        let pairMap =
                            if map.ContainsKey watchTime then
                                let temp = map.[watchTime]
                                let watchList =
                                    if temp.ContainsKey postEvent then
                                        preEvent :: temp.[postEvent]
                                    else
                                        [ preEvent ]
                                Map.add postEvent watchList temp
                            else
                                Map.add postEvent [ preEvent ] Map.empty
                        Map.add watchTime pairMap map
                    ) currentWatchMap
            else
                currentWatchMap

        let CheckTriggers postEvent (currentWatchMap:Map<int, Map<int * int, (int * int * EventLabel) list>>) (currentResultList) =
            let time, nindex, _ = postEvent

            seq {
                for t in (time-jitter)..(time+jitter) -> t
                }
            |> Seq.filter (fun t -> currentWatchMap.ContainsKey t)
            |> Seq.fold (fun (state:Map<int, Map<int * int, (int * int * EventLabel) list>> * (int * int * EventLabel) list) t ->
                    let watchMap, resultList = state
                    let pairMap = watchMap.[t]
                    let eventKey = t, nindex
                    let updatedMap, updatedResultList =
                        if pairMap.ContainsKey eventKey then
                            let watchSet =
                                Set.add postEvent (Set.ofList pairMap.[eventKey])
                            (Map.add time (Map.remove eventKey pairMap) watchMap), (List.append (Set.toList watchSet) resultList)
                        else
                            watchMap, resultList
                    updatedMap, updatedResultList
                ) (currentWatchMap, currentResultList)

        // watchMap: Key = time; Value = Map<postEvent, preEvent list>
        let _, selectedFiringEvents =
            firingEvents
            |> Seq.fold (fun (state:(Map<int, Map<int * int, (int * int * EventLabel) list>>) * ((int * int * EventLabel) list)) event ->
                    let watchMap, resultList = state

                    // remove old watch events
                    let updatedMap1 =
                        if (resultList.Length % 100) = 0 then
                            CleanMap event jitter watchMap
                        else
                            watchMap

                    let updatedMap2 = SelectWatchedEvents mapPrePost event updatedMap1
                    CheckTriggers event updatedMap2 resultList
                ) (Map.empty, [])

        selectedFiringEvents

    // Generate a MapPrePost from the network structure suitable for passing to FilterEvents
    let GetPairDataFromNetworkStructure (spn:CrossbarNetwork) minWeight =
        spn.Connections
        |> Seq.filter (fun connection -> connection.Weight > minWeight)
        |> Seq.map (fun connection -> connection.PreNeuron, connection.PostNeuron, int(connection.Delay))
        |> Seq.fold (fun (state:Map<int, _>) (pre, post, delay) ->
                let postList =
                    if state.ContainsKey pre then
                        state.[pre]
                    else
                        []
                Map.add pre ((post, delay, 0) :: postList) state
            ) Map.empty

    // Given a triggering event (e.g. one of a triplet), find the longest path through the network of firing events
    // The longest path is usually discoverable from any one of the trigger events
    // except on rare occasions where all triggers must be tried
    // the result for the template should match the max link level returned from the Izhikevich search algorithm
    // The prePostEventMap is a map of paired firing events from a PNGDescriptor or LinkedEventPair list
    // i.e. for each pre event, map the list of post events
    let GetLongestEventPath triggerEvent prePostEventMap =

        let rec FindLongestPath triggerEvent (prePostEventMap:Map<_, _>) (links:(int * int) list) =
            if not (prePostEventMap.ContainsKey triggerEvent) then
                links
            else
                prePostEventMap.[triggerEvent]
                |> List.map (fun ev -> FindLongestPath ev prePostEventMap (ev :: links))
                |> Seq.maxBy (fun links -> links.Length)
        
        FindLongestPath triggerEvent prePostEventMap [triggerEvent]
